package com.langll.ai;

import com.langll.ai.tool.DateTimeTools;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbacks;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 *  测试tools
 */
@SpringBootTest(classes = AIChatApplication.class)
@RunWith(SpringRunner.class)
public class TestDateTimeTools {

    @Autowired
    private ChatClient chatClient;
    @Test
    public void generateStream(){

        ToolCallback[] dateTimeTools = ToolCallbacks.from(new DateTimeTools());
        ChatOptions chatOptions = ToolCallingChatOptions.builder()
                .toolCallbacks(dateTimeTools)
                .build();

        Prompt prompt = new Prompt("今天是几号?", chatOptions);

        ChatClient.CallResponseSpec responseSpec = chatClient.prompt(prompt).call();
        System.out.println(responseSpec.content());
    }

    @Test
    public void updateCustomerInfo(){

        ToolCallback[] dateTimeTools = ToolCallbacks.from(new DateTimeTools());
        ChatOptions chatOptions = ToolCallingChatOptions.builder()
                .toolCallbacks(dateTimeTools)
                .build();

        Prompt prompt = new Prompt("将id是100的客户，名称更新成张三，邮箱更新成zhangsan@163.com", chatOptions);

        ChatClient.CallResponseSpec responseSpec = chatClient.prompt(prompt).call();
        System.out.println(responseSpec.content());
    }

    @Test
    public void getCustomerInfo(){
        ToolCallback[] dateTimeTools = ToolCallbacks.from(new DateTimeTools());
        ChatOptions chatOptions = ToolCallingChatOptions.builder()
                .toolCallbacks(dateTimeTools)
                .toolContext(Map.of("tenantId", "acme"))
                .build();

        Prompt prompt = new Prompt("检索客户信息id为100", chatOptions);

        ChatClient.CallResponseSpec responseSpec = chatClient.prompt(prompt).call();
        System.out.println(responseSpec.content());
    }

    @Test
    public void getCustomerInfo2(){
        List<Message> messages = new ArrayList<>();
        UserMessage userMessage = new UserMessage("检索客户信息id为100");
        messages.add(userMessage);
        Prompt prompt = new Prompt(messages);

        ChatClient.CallResponseSpec response = chatClient.prompt(prompt)
                .tools("getCustomerInfo")
                .toolContext(Map.of("tenantId", "acme"))
                .call();

        System.out.println(response.content());
    }
}
